# !/usr/bin/env python
# -*- coding=utf-8 -*-


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as md
import warnings
from main import get_data
from pyecharts import options as opts
from pyecharts.charts import Kline

warnings.filterwarnings("ignore")
plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
pd.options.display.max_rows = 20


BAR_DATA_PATH_PX_MV = r"C:\Users\新田草\Desktop\正仁量化\data\bar_data_px_mv.csv"
BAR_DATA_PATH_TR = r"C:\Users\新田草\Desktop\正仁量化\data\bar_data_tr.csv"
FIGURE_SAVE_PATH = r"C:\Users\新田草\Desktop\正仁量化\candle_figure"


data = get_data(BAR_DATA_PATH_TR)
data = data.iloc[:, 1:]

# 画length根 K 线
length = 100
dates, open_price, high_price, \
low_price, close_price, volumes = data["time"][:length], \
                                  data["open"][:length], \
                                  data["high"][:length], \
                                  data["low"][:length], \
                                  data["close"][:length], \
                                  data["volume"][:length]


def plot_candle_with_pyecharts():
    y_data = [[o, h, l, c] for o, h, l, c in
              zip(open_price, high_price, low_price, close_price)]
    x_data = list(dates)

    c = (
     Kline()
     .add_xaxis(xaxis_data=x_data)
     .add_yaxis(
      "CY1805",
      y_axis=y_data,
      itemstyle_opts=opts.ItemStyleOpts(
       color="#ec0000",
       color0="#00da3c",
       border_color="#8A0000",
       border_color0="#008F28",
      ),
     )
     .set_global_opts(
      xaxis_opts=opts.AxisOpts(is_scale=True),
      yaxis_opts=opts.AxisOpts(
       is_scale=True,
       splitarea_opts=opts.SplitAreaOpts(
        is_show=True, areastyle_opts=opts.AreaStyleOpts(opacity=1)
       ),
      ),
      datazoom_opts=[opts.DataZoomOpts(type_="inside")],
      title_opts=opts.TitleOpts(title="Kline-ItemStyle"),
     )
     .render("K线图鼠标缩放.html")
    )


def plot_candle_with_matplotlib():
    # 2.设置绘图窗口
    plt.figure(figsize=(12, 8), facecolor="lightgray")
    plt.title("CY1805", fontsize=16)
    plt.xlabel("Data", fontsize=14)
    plt.ylabel("Price", fontsize=14)

    contract_multiplier = 5
    up = high_price.max() + contract_multiplier*10
    bottom = low_price.min() - contract_multiplier*10

    plt.ylim([bottom, up])

    # 3.x坐标（时间轴）轴修改
    ax = plt.gca()
    # y 轴不可见
    # ax.axes.get_yaxis().set_visible(False)
    # x 轴不可见
    ax.axes.get_xaxis().set_visible(False)

    # 4.判断收盘价与开盘价 确定蜡烛颜色
    colors_bool = close_price >= open_price
    colors = np.zeros(colors_bool.size, dtype="U5")
    colors[:] = "green"
    colors[colors_bool] = "red"

    # 5.确定蜡烛边框颜色
    edge_colors = np.zeros(colors_bool.size, dtype="U1")
    edge_colors[:] = "g"
    edge_colors[colors_bool] = "r"

    # 6.绘制蜡烛
    plt.bar(dates, (close_price - open_price), 0.8, bottom=open_price, color=colors,
           edgecolor=edge_colors, zorder=3)

    # 7.绘制蜡烛直线(最高价与最低价)
    plt.vlines(dates, low_price, high_price, color=edge_colors)
    plt.savefig(FIGURE_SAVE_PATH+"/candle_fig.png", figsize=[10, 8])
    plt.show()


if __name__ == '__main__':
    # 生成matplotlib风格的 K 线图
    plot_candle_with_matplotlib()
    # 生成pyecharts风格可动态缩放的 K 线图
    plot_candle_with_pyecharts()

